import autogen

# Define your LLM config
llm_config = {
    "config_list": [
        {
            "model": "gpt-4o",
            "api_key": "sk-...",  # or use env var
        }
    ],
    "temperature": 0.1,      # low = more deterministic
    "cache_seed": 42,        # reproducible runs (optional)
}

# Or load from a JSON file (recommended for production)
config_list = autogen.config_list_from_json(
    "OAI_CONFIG_LIST"
)

# Termination methods:

# 1. Max turns
user_proxy.initiate_chat(assistant, max_turns=5)

# 2. Keyword in reply ("TERMINATE")
assistant = autogen.AssistantAgent(
    system_message="""...(your instructions)...
    When done, reply with: TERMINATE"""
)

# 3. Custom function
def my_termination(msg):
    return "task_complete" in msg["content"].lower()

user_proxy = autogen.UserProxyAgent(
    is_termination_msg=my_termination
)

# Install AutoGen (v0.4+)
pip install pyautogen

# For code execution in Docker (recommended)
pip install pyautogen[docker]

# Set your API key
export OPENAI_API_KEY="sk-..."

import autogen

# 1. LLM config
llm_config = {
    "config_list": [{"model": "gpt-4o", "api_key": "sk-..."}]
}

# 2. The AI assistant agent
assistant = autogen.AssistantAgent(
    name="assistant",
    llm_config=llm_config,
    system_message="""You are a helpful Python expert.
    When you finish the task, reply with: TERMINATE"""
)

# 3. The user proxy (executes code, no human input)
user_proxy = autogen.UserProxyAgent(
    name="user_proxy",
    human_input_mode="NEVER",
    is_termination_msg=lambda x: "TERMINATE" in x.get("content", ""),
    code_execution_config={
        "work_dir": "coding",
        "use_docker": False,  # set True in production
    }
)

# 4. Start the conversation!
user_proxy.initiate_chat(
    assistant,
    message="Write a Python script that prints the first 10 Fibonacci numbers."
)

# Step 1: Writer agent creates a blog post
result1 = user_proxy.initiate_chat(
    writer, message="Write a blog post about RAG systems"
)
draft = result1.summary

# Step 2: Critic reviews it
result2 = user_proxy.initiate_chat(
    critic, message=f"Review this blog post:\n\n{draft}"
)

# Step 3: Editor applies improvements
result3 = user_proxy.initiate_chat(
    editor, message=f"Apply this feedback:\n{result2.summary}"
)

# Register a nested chat — when assistant gets a coding task,
# it spawns a full sub-conversation with a coding specialist
assistant.register_nested_chats(
    trigger=user_proxy,
    chat_queue=[
        {
            "recipient": coding_specialist,
            "message": "Please implement this: ",
            "summary_method": "last_msg",
            "max_turns": 3,
        }
    ]
)

from autogen import SwarmAgent, initiate_swarm_chat

# Each agent defines who it can hand off to
triage = SwarmAgent(
    name="triage",
    handoffs=["billing_agent", "tech_support", "sales"]
)

initiate_swarm_chat(
    initial_agent=triage,
    agents=[triage, billing, tech_support, sales],
    messages="I can't access my account after payment failed"
)

import autogen
from autogen import AssistantAgent, UserProxyAgent

llm_config = {"config_list": [{"model": "gpt-4o", "api_key": "..."}]}

assistant = AssistantAgent(name="assistant", llm_config=llm_config)
user_proxy = UserProxyAgent(
    name="user_proxy", human_input_mode="NEVER",
    code_execution_config=False  # disable code exec, use tools instead
)

# ✨ Register a tool — function is callable by user_proxy,
#    description is shown to the LLM to know when to call it
@user_proxy.register_for_execution()
@assistant.register_for_llm(description="Get current weather for a city")
def get_weather(city: str) -> str:
    # Call a real weather API here
    return f"Weather in {city}: 22°C, sunny"

@user_proxy.register_for_execution()
@assistant.register_for_llm(description="Search the web for current info")
def web_search(query: str) -> str:
    # Integrate with Serper, Tavily, or Bing here
    return f"Results for '{query}': ..."

user_proxy.initiate_chat(
    assistant,
    message="What's the weather in Tokyo and is there a tech conference there this week?"
)

import autogen

llm_config = {"config_list": [{"model": "gpt-4o", "api_key": "..."}]}

# Specialized agents with distinct system messages
planner = autogen.AssistantAgent(
    name="Planner", llm_config=llm_config,
    system_message="You break complex tasks into subtasks and assign them."
)

coder = autogen.AssistantAgent(
    name="Coder", llm_config=llm_config,
    system_message="You write high-quality Python code. No prose, just code."
)

critic = autogen.AssistantAgent(
    name="Critic", llm_config=llm_config,
    system_message="Review code for bugs, edge cases, and style. Be concise."
)

user_proxy = autogen.UserProxyAgent(
    name="User", human_input_mode="NEVER",
    code_execution_config={"work_dir": "coding"}
)

# Assemble the group chat
groupchat = autogen.GroupChat(
    agents=[user_proxy, planner, coder, critic],
    messages=[],
    max_round=12,
    speaker_selection_method="auto"  # LLM picks next speaker
)

# The manager orchestrates the conversation
manager = autogen.GroupChatManager(
    groupchat=groupchat,
    llm_config=llm_config
)

user_proxy.initiate_chat(
    manager,
    message="Build a FastAPI endpoint that analyzes sentiment in text"
)

def custom_speaker_selector(last_speaker, groupchat):
    # Always follow coder with critic
    if last_speaker.name == "Coder":
        return next(a for a in groupchat.agents if a.name == "Critic")
    # Always follow critic with user_proxy (to run fixed code)
    if last_speaker.name == "Critic":
        return next(a for a in groupchat.agents if a.name == "User")
    return "auto"  # let LLM decide otherwise

groupchat = autogen.GroupChat(
    ...,
    speaker_selection_method=custom_speaker_selector
)

from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent

rag_agent = RetrieveUserProxyAgent(
    name="rag_agent",
    retrieve_config={
        "task": "qa",
        "docs_path": ["./my_docs/", "https://example.com/api-docs"],
        "chunk_token_size": 2000,
        "model": "gpt-4o",
        "vector_db": "chroma",   # or "pgvector", "qdrant"
        "collection_name": "my_docs",
        "get_or_create": True,     # reuse existing collection
    },
    code_execution_config=False,
    human_input_mode="NEVER"
)

rag_agent.initiate_chat(
    assistant,
    problem="What does our API return when authentication fails?"
)

import chromadb

chroma_client = chromadb.Client()
memory_collection = chroma_client.get_or_create_collection("agent_memory")

# Give agents tools to read/write memory
@user_proxy.register_for_execution()
@assistant.register_for_llm(description="Save a fact to long-term memory")
def save_memory(key: str, value: str) -> str:
    memory_collection.add(
        documents=[value],
        ids=[key]
    )
    return f"Saved: {key}"

@user_proxy.register_for_execution()
@assistant.register_for_llm(description="Search memory for relevant facts")
def search_memory(query: str) -> str:
    results = memory_collection.query(query_texts=[query], n_results=3)
    return str(results["documents"])